package com.example.rpc.demo.registry;

import com.example.rpc.demo.consts.RpcConstant;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.concurrent.CountDownLatch;


@Component
public class ServiceRegistry {
    private static final Logger logger = LoggerFactory.getLogger(ServiceRegistry.class);
    private ZooKeeper zk = null;

    private CountDownLatch latch = new CountDownLatch(1);

    @Value("${zk.registry.address}")
    private String registryAddress;

    /**
     * 将服务名与地址进行注册
     * @param serviceName
     * @param address
     */
    public void registry(String serviceName,String address) {
        if (zk == null) {
            zk = connectServer();
        }

        if(zk!=null) {
            String path = createServicePath(zk, serviceName);
            if(path ==null) {
                return;
            }

            createAddressNode(zk, path, address);
            logger.info("registry over .");
        } else {
            logger.error("zk is null.");
        }
    }

    private ZooKeeper connectServer() {
        ZooKeeper zk = null;
        try {
            zk = new ZooKeeper(this.registryAddress, RpcConstant.ZK_SESSION_TIMEOUT, watchedEvent -> {
                if(watchedEvent.getState() == Watcher.Event.KeeperState.SyncConnected) {
                    latch.countDown();
                }
            });
            latch.await();
        } catch (IOException | InterruptedException e) {
            logger.error("connect zk server failed:", e);
        }
        logger.info("connect zk success");

        return zk;
    }

    private String createServicePath(ZooKeeper zk,String serviceName) {
        try {
            Stat stat = zk.exists(RpcConstant.ZK_REGISTRY_PATH + "/" + serviceName, false);
            if(stat ==null) {
                String path = zk.create(RpcConstant.ZK_REGISTRY_PATH + "/" + serviceName, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                logger.info("create zookeeper path({}=>{})",serviceName,path);
                return path;
            }

            return RpcConstant.ZK_REGISTRY_PATH+"/"+serviceName;

        } catch (KeeperException | InterruptedException e) {
            logger.error("create zk node failed: ", e);
            return null;
        }
    }

    private void createAddressNode(ZooKeeper zk,String path,String address) {
        try {
            byte[] bytes = address.getBytes();

            String node = zk.create(path+"/", bytes, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            logger.info("create zookeeper node ({} => {})", address, node);
        } catch (KeeperException | InterruptedException e) {
            logger.error("create zk node failed .", e);
        }
    }


}
